#include "llvm/ADT/StringExtras.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/Pass.h"
#include "llvm/Passes/PassBuilder.h"
#include "llvm/Passes/PassPlugin.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/raw_ostream.h"
#include <limits>
#include <random>

using namespace llvm;

static cl::opt<bool> Wave("wave-custom-ollvm", cl::init(false),
                          cl::desc("wave custom ollvm"));

namespace {

bool runCustomOllvm(Function &F) {
  if (Wave) {
    errs() << "Custom Ollvm: ";
    errs().write_escaped(F.getName()) << '\n';
  }
  return false;
}

struct LegacyCustomOllvm : public FunctionPass {
  static char ID;
  LegacyCustomOllvm() : FunctionPass(ID) {}
  bool runOnFunction(Function &F) override { return runCustomOllvm(F); }
};

struct CustomOllvm : PassInfoMixin<CustomOllvm> {
  CustomOllvm();

  PreservedAnalyses run(Function &F, FunctionAnalysisManager &) {
    if (F.getName().compare("main") != 0) {
      // 构造新名字

      std::string NewFunctionName;
      genNewFunctionName(F, NewFunctionName);
      F.setName(NewFunctionName);
    }

    return PreservedAnalyses::none();
  }

  void genNewFunctionName(Function &F, std::string& NewFunctionName);

  static bool isRequired() { return true; }

private:
  int FunctionRandomId = 0;
};

} // namespace

CustomOllvm::CustomOllvm() {
  std::random_device Rd;  // 用于获取随机种子
  std::mt19937 Gen(Rd()); // 使用随机种子初始化一个随机数生成器
  // 获取int的最大值
  int MaxInt = std::numeric_limits<int>::max();
  // 获取int的最小值
  int MinInt = std::numeric_limits<int>::min();
  std::uniform_int_distribution<> Dis(MinInt, MaxInt); // 定义一个在1到100之间的均匀分布
  FunctionRandomId = Dis(Gen); // 生成一个随机数
}

char LegacyCustomOllvm::ID = 0;

static RegisterPass<LegacyCustomOllvm> X("customOllvm", "CustomOllvm Pass",
                                 false /* Only looks at CFG */,
                                 false /* Analysis Pass */);

/* New PM Registration */
llvm::PassPluginLibraryInfo getCustomOllvmPluginInfo() {
  return {LLVM_PLUGIN_API_VERSION, "CustomOllvm", LLVM_VERSION_STRING,
          [](PassBuilder &PB) {
            PB.registerVectorizerStartEPCallback(
                [](llvm::FunctionPassManager &PM, OptimizationLevel Level) {
                  PM.addPass(CustomOllvm());
                });
            PB.registerPipelineParsingCallback(
                [](StringRef Name, llvm::FunctionPassManager &PM,
                   ArrayRef<llvm::PassBuilder::PipelineElement>) {
                  if (Name == "customOllvm") {
                    PM.addPass(CustomOllvm());
                    return true;
                  }
                  return false;
                });
          }};
}

void CustomOllvm::genNewFunctionName(
    Function &F, std::string& NewFunctionName) {
  // 获取当前函数名
  StringRef OldName = F.getName();

  // 构造新名字（例如，在原名字前加前缀 "renamed_"）
  std::string NewName = ("renamed_" + OldName).str();
  NewName += FunctionRandomId;

  llvm::MD5 Hasher;
  llvm::MD5::MD5Result Hash;
  Hasher.update(NewName);
  Hasher.final(Hash);
  NewFunctionName = "123_" + llvm::utohexstr(Hash.low(), /*LowerCase=*/true);
}

#ifndef LLVM_CUSTOM_OLLVM_LINK_INTO_TOOLS
extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo
llvmGetPassPluginInfo() {
  return getCustomOllvmPluginInfo();
}
#endif
